import numpy as np
from matplotlib import pyplot as plt
from sklearn.preprocessing import StandardScaler
from tensorflow import keras
from tensorflow.python.keras.layers import Dense
from tensorflow.python.keras.models import Sequential
import tensorflow as tf
import getData

seed = 42
np.random.seed(seed)
tf.random.set_seed(seed)
epochs = 1000


class Fine_tune():
    def fit(self, data):
        self.x_source_train = data.source_x_train
        self.y_source_train = data.source_y_train

        self.x_target_train = data.target_x_train
        self.y_target_train = data.target_y_train

        self.init_data()
        self.model_build()
        self.model_train()
        self.model_transfer()
        self.fine_tune()

    def init_data(self):
        # 验证集
        # source
        self.x_source_valid = self.x_source_train[3:self.x_source_train.shape[0]: 5]
        x_train1 = self.x_source_train[0:self.x_source_train.shape[0]:5]
        x_train2 = self.x_source_train[1:self.x_source_train.shape[0]:5]
        x_train3 = self.x_source_train[2:self.x_source_train.shape[0]:5]
        x_train4 = self.x_source_train[4:self.x_source_train.shape[0]:5]
        self.x_source_train_ori = np.concatenate((x_train1, x_train2, x_train3, x_train4))

        self.y_source_valid = self.y_source_train[3:self.y_source_train.shape[0]:5]
        y_train1 = self.y_source_train[0:self.y_source_train.shape[0]:5]
        y_train2 = self.y_source_train[1:self.y_source_train.shape[0]:5]
        y_train3 = self.y_source_train[2:self.y_source_train.shape[0]:5]
        y_train4 = self.y_source_train[4:self.y_source_train.shape[0]:5]
        self.y_source_train_ori = np.concatenate((y_train1, y_train2, y_train3, y_train4))
        # target
        self.x_target_valid = self.x_target_train[3:self.x_target_train.shape[0]: 5]
        x_train1 = self.x_target_train[0:self.x_target_train.shape[0]:5]
        x_train2 = self.x_target_train[1:self.x_target_train.shape[0]:5]
        x_train3 = self.x_target_train[2:self.x_target_train.shape[0]:5]
        x_train4 = self.x_target_train[4:self.x_target_train.shape[0]:5]
        self.x_target_train_ori = np.concatenate((x_train1, x_train2, x_train3, x_train4))

        self.y_target_valid = self.y_target_train[3:self.y_target_train.shape[0]:5]
        y_train1 = self.y_target_train[0:self.y_target_train.shape[0]:5]
        y_train2 = self.y_target_train[1:self.y_target_train.shape[0]:5]
        y_train3 = self.y_target_train[2:self.y_target_train.shape[0]:5]
        y_train4 = self.y_target_train[4:self.y_target_train.shape[0]:5]
        self.y_target_train_ori = np.concatenate((y_train1, y_train2, y_train3, y_train4))

        # 归一化
        self.x_scaler = StandardScaler().fit(self.x_source_train)
        self.y_scaler = StandardScaler().fit(self.y_source_train)

        self.x_source_train_st = self.x_scaler.transform(self.x_source_train_ori)
        self.y_source_train_st = self.y_scaler.transform(self.y_source_train_ori)

        self.x_target_train_st = self.x_scaler.transform(self.x_target_train_ori)
        self.y_target_train_st = self.y_scaler.transform(self.y_target_train_ori)

        self.x_source_valid = self.x_scaler.transform(self.x_source_valid)
        self.y_source_valid = self.y_scaler.transform(self.y_source_valid)

        self.x_target_valid = self.x_scaler.transform(self.x_target_valid)
        self.y_target_valid = self.y_scaler.transform(self.y_target_valid)

    def model_build(self):
        model = Sequential()
        model.add(Dense(13, input_dim=1, activation='relu', name="layer2"))
        model.add(Dense(13, activation='relu', name="layer3"))
        model.add(Dense(2, name="out"))
        self.model = model

    def model_train(self):
        self.model.compile(optimizer=keras.optimizers.Nadam(), loss="MSE")
        self.history = self.model.fit(self.x_source_train_st, self.y_source_train_st, epochs=epochs,
                                      validation_data=(self.x_source_valid, self.y_source_valid),
                                      batch_size=8)  # fit开始训练

    def model_transfer(self):
        for layer in self.model.layers[:1]:
            layer.trainable = False
        print(self.model.summary())
        self.model.compile(optimizer=keras.optimizers.Nadam(), loss="MSE")
        self.history = self.model.fit(self.x_target_train_st, self.y_target_train_st, epochs=epochs,
                                      validation_data=(self.x_target_valid, self.y_target_valid),
                                      batch_size=8)  # fit开始训练

    def fine_tune(self):
        # 解冻
        for layer in self.model.layers:
            layer.trainable = True
        self.model.compile(optimizer=keras.optimizers.Nadam(1e-3), loss="MSE")
        self.history = self.model.fit(self.x_target_train_st, self.y_target_train_st, epochs=epochs,
                                      validation_data=(self.x_target_valid, self.y_target_valid),
                                      batch_size=8)  # fit开始训练

    def predict(self, x_test):
        # 归一化
        self.x_test = self.x_scaler.transform(x_test)
        y_pre = self.model.predict(self.x_test)
        y_pre = self.y_scaler.inverse_transform(y_pre)
        return y_pre


if __name__ == '__main__':
    model = Fine_tune()
    data = getData.GetData()
    data.getData()
    model.fit(data)
    x_test = data.target_x_test
    y_pre = model.predict(x_test)

    y = data.target_y_test
    for i in range(y_pre.shape[1]):
        plt.figure()
        plt.plot(x_test, y_pre[:, i], color="red")
        plt.plot(x_test, y[:, i])
        plt.show()
